using Distributed
addprocs(5)

@everywhere using SharedArrays
@everywhere using Parameters, Interpolations, Plots, Distributions, Random, CSV, DataFrames, NLopt, Optim
@everywhere include("frhc_setup.jl")
@everywhere include("frhc_readin.jl")
@everywhere include("frhc_valfunc.jl")
@everywhere include("frhc_simulate.jl")
@everywhere include("frhc_estimate.jl")

@everywhere dir = "C:\\Users\\Garrett\\Documents\\Work\\Papers\\FRHC\\replication_archive\\Model"
@everywhere utilities, moments = Readin_utilities(dir), Readin_moments(dir)
@everywhere cd(dir)

#starting guess
tempthing = CSV.read("$dir\\estimates_20241018.csv", DataFrame; skipto=1)
tempthing2 = Array{Float64,2}(tempthing)'
tempthing3 = zeros(length(tempthing2))
for iter = 1:length(tempthing2)
    tempthing3[iter] = tempthing2[iter]
end
guess_init = copy(tempthing3)
println(round.(guess_init, digits=4))

deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500)

error_estim = deliverables[1]
error_estim

CSV.write("simulated_data\\simulated_data_base.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

##################Recession##########################
deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=1, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_rec.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

#no parent PBR
deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=0, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_rec_nopbr.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file


####################Borrowing Constraint Modifications##########################
#No borrowing constraints
deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, dlim = 4, pbr = 1)
CSV.write("simulated_data\\simulated_data_base_nolim.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

####various borrowing constraints in recession
for i = 1:4
    deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=1, rec_pol = 1, dlim=i)
    CSV.write("simulated_data\\simulated_data_rec_dlim_$i.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file    
end

#############Recession decomposition
for d = 1:3
    deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, rec_pol = 1, pbr=1, decomp = d)
    CSV.write("simulated_data\\simulated_data_rec_decomp_$d.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file
end

##############Loop Over Counterfactuals##########################
#three Pell counterfactuals: expansion with progressivity elimination, expanded generosity, and wealth consideration
#run in/out of recession and with/without parental behavioral response
for i_c = 1:3, i_p = 0:1, i_r=0:1
    deliverables = Objective_function(guess_init, utilities, moments; ret = 1, cfact = i_c, pbr=i_p, rec=i_r, rec_pol = i_r, nsim = 2500)
    filename = "simulated_data_cfact_" * "$i_c" * "_pbr_" * "$i_p" * "_rec_" * "$i_r"
    CSV.write("simulated_data\\$filename.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file
end

########Recession policy######

#no policy in recession world
deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=1, rec_pol = 0)
CSV.write("simulated_data\\simulated_data_rec_nopol.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

#implementing in various scenarios
deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 0, pbr=0, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_cfact_4_pbr_0_rec_0.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 0, pbr=1, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_cfact_4_pbr_1_rec_0.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=0, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_cfact_4_pbr_0_rec_1.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file

deliverables = Objective_function(guess_init, utilities, moments; ret = 1, nsim = 2500, rec = 1, pbr=1, rec_pol = 1)
CSV.write("simulated_data\\simulated_data_cfact_4_pbr_1_rec_1.csv", DataFrame(deliverables[3], :auto), header=false) #write CSV file








##########################Standard Error Function##########################
using LinearAlgebra

#function for computing standard errors
function standard_errors(guess::Vector{Float64}, perturb::Float64)
    num_sim = 2500 #10x normal simulation
    nrep = 10
    num_param = length(guess)
    num_moment = 582
    perturb_mat = zeros(num_moment, num_param)
    se_vec = zeros(num_param)

    #setup
    utilities = Readin_utilities(dir)
    moments = Readin_moments(dir)
    err_vec = Objective_function(guess, utilities, moments; nsim = num_sim, ret = 1)[2] #
    
    #loop over parameters
    for i = 1:num_param
        guess_perturb = copy(guess)
        guess_perturb[i] += perturb
        err_vec_perturb = Objective_function(guess_perturb, utilities, moments; nsim = num_sim, ret = 1)[2] #
        perturb_mat[:, i] = (err_vec.-err_vec_perturb)./perturb
    end

    gradient_mat = perturb_mat
    W = Diagonal(ones(num_moment))
    #W[526, 526] = 25.0
    #W[527, 527] = 25.0
    #W[528, 528] = 25.0
    #W[529, 529] = 4.0
    #W[530, 530] = 25.0
    avar_mat_step1 = Matrix{Float64}(gradient_mat' * W * gradient_mat)
    avar_mat_step2 = pinv(avar_mat_step1)
    avar_mat_step3 = (1 + 1/nrep).*avar_mat_step2
    se_vec = sqrt.(diag(avar_mat_step3))
    se_vec
end

#get standard erors
tempthing = CSV.read("$dir\\estimates_20241018.csv", DataFrame; skipto=1)
tempthing2 = Array{Float64,2}(tempthing)'
tempthing3 = zeros(length(tempthing2))
for iter = 1:length(tempthing2)
    tempthing3[iter] = tempthing2[iter]
end
guess_init = copy(tempthing3)
println(round.(guess_init, digits=4))

se_vec = standard_errors(guess_init, 0.005)
println(se_vec)
CSV.write("$dir\\ses_20241018.csv", DataFrame(se_vec', :auto), header=false) #write CSV file


####end of file
